from cnn.zstd import decompress_file
from io import BytesIO
import pandas as pd
import torch
from cnn.func import LoadData, normalization
from torch.utils.data import DataLoader


from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.preprocessing import LabelEncoder


# Pands导入数据
def preprocess(input_file):
    a = decompress_file(input_file)
    df = pd.read_csv(BytesIO(a))

    df.columns = [
        "duration",
        "protocol_type",
        "service",
        "flag",
        "src_bytes",
        "dst_bytes",
        "land",
        "wrong_fragment",
        "urgent",
        "hot",
        "num_failed_logins",
        "logged_in",
        "num_compromised",
        "root_shell",
        "su_attempted",
        "num_root",
        "num_file_creations",
        "num_shells",
        "num_access_files",
        "num_outbound_cmds",
        "is_host_login",
        "is_guest_login",
        "count",
        "srv_count",
        "serror_rate",
        "srv_serror_rate",
        "rerror_rate",
        "srv_rerror_rate",
        "same_srv_rate",
        "diff_srv_rate",
        "srv_diff_host_rate",
        "dst_host_count",
        "dst_host_srv_count",
        "dst_host_same_srv_rate",
        "dst_host_diff_srv_rate",
        "dst_host_same_src_port_rate",
        "dst_host_srv_diff_host_rate",
        "dst_host_serror_rate",
        "dst_host_srv_serror_rate",
        "dst_host_rerror_rate",
        "dst_host_srv_rerror_rate",
        "label",
    ]

    # one-hot编码
    # 数值列
    number_col = df.select_dtypes(include=["number"]).columns

    # 分类变量
    cat_col = df.columns.difference(number_col)

    cat_col = cat_col.drop("label")

    # 将分类变量筛选出来
    df_cat = df[cat_col].copy()

    # one-hot编码
    one_hot_data = pd.get_dummies(df_cat, columns=cat_col)

    # 将原数据的分类变量去掉
    one_hot_df = pd.concat([df, one_hot_data], axis=1)
    one_hot_df.drop(columns=cat_col, inplace=True)

    # 归一化
    minmax_scale = MinMaxScaler(feature_range=(0, 1))

    normalized_df = normalization(one_hot_df.copy(), number_col, minmax_scale)
    # 标签编码
    # 为不同的类别进行编码
    labels = pd.DataFrame(df.label)
    label_encoder = LabelEncoder()
    enc_label = labels.apply(label_encoder.fit_transform)
    normalized_df.label = enc_label
    label_encoder.classes_
    data = normalized_df
    # 定义训练参数
    X = data.drop(columns=["label"])

    y = data["label"]
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.20, random_state=50
    )

    train_data = LoadData(X_train, y_train)
    test_data = LoadData(X_test, y_test)
    X_dimension = len(X_train.columns)
    y_dimension = len(y_train.value_counts())

    print(f"X的维度：{X_dimension}")
    print(f"y的维度：{y_dimension}")
    batch_size = 128

    train_dataloader = DataLoader(train_data, batch_size=batch_size)
    test_dataloader = DataLoader(test_data, batch_size=batch_size)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    return (
        device,
        train_dataloader,
        test_dataloader,
        train_data,
        test_data,
        X_dimension,
        y_dimension,
    )
